#ifndef PCPS_CONJUGATE_GRADIENT_HEADER
#define PCPS_CONJUGATE_GRADIENT_HEADER

#include <src/pcps.h>
#include <src/matrix_actor.h>

namespace pcps {

  //-----------------------------------------------------------------------------------
  // pcps::conjugate_gradient -- Solves a system of linear equations A x = b
  //                             using the conjugate gradient method.
  //
  //                             Using preconditioner M, the equations become
  //
  //                             (M A M) v = (M b)           v = M^(-1) x
  //
  //                             which is solved for v by conjugate gradient.
  //
  //                             x is then computed as x = M v.
  //
  //                             A must be a positive definite hermitian matrix.
  //                             M must be an invertable hermitian matrix.
  //
  //                             Presumably, the ratio of largest to smallest
  //                             eigenvalues is smaller for (M A M) than for A.
  //-----------------------------------------------------------------------------------

  template<class S>
  inline void conjugate_gradient(const int n,                     // dimension of the system
                                 const int m,                     // maximum number of iterations
                                 S * const x,                // size n.                             on exit, the solution vector x.
                                 S * const b,                // size n.  on entry, the vector b.    on exit, contents are overwritten.
                                 S * const y,                // size n.  vector used as workspace.  on exit, contents are overwritten.
                                 S * const d,                // size n.  vector used as workspace.  on exit, contents are overwritten.
                                 S * const r,                // size n.  vector used as workspace.  on exit, contents are overwritten.
                                 pcps::MatrixActor<S> * a,   // object that applies linear transformations
                                 const int tid = 0)               // thread id
  {

    // get MPI info
    const MPI::Comm & comm = MPI::COMM_WORLD;
    const int nproc = comm.Get_size();
    const int myrank = comm.Get_rank();

    S r_norm = pcps::unity<S>() * 1.0e100;
    S alpha, beta;

    // initialize the residual as r = M b
    a->operate_by_M(n, b, r, myrank, tid);

    // initialize the search direction as d = r
    if (tid == 0) pcps::xcopy(n, r, 1, d, 1);

    // use the b vector to hold v
    S * const v = b;

    // initialize v to zero
    if (tid == 0) pcps::xscal(n, pcps::zero<S>(), v, 1);

    // compute the initial residual norm
    if (tid == 0) r_norm = pcps::xdotc(n, r, 1, r, 1);

    // iteratively compute v
    int iter = 0;
    for ( ; iter < m; iter++) {

      // check if we are finished 
      const bool converged = a->converged(std::abs(std::sqrt(r_norm)), myrank, tid);
      if (converged) break;

      // compute y = (M A M) d       (use x as an intermediate)
      a->operate_by_M(n, d, y, myrank, tid);
      a->operate_by_A(n, y, x, myrank, tid);
      a->operate_by_M(n, x, y, myrank, tid);

      // compute alpha
      if (tid == 0) alpha = r_norm / pcps::xdotc(n, d, 1, y, 1);

      // update v
      if (tid == 0) pcps::xaxpy(n, alpha, d, 1, v, 1);

      // update the residual
      if (tid == 0) pcps::xaxpy(n, -alpha, y, 1, r, 1);

      // compute beta
      if (tid == 0) beta = pcps::xdotc(n, r, 1, r, 1) / r_norm;

      // compute new residual norm
      if (tid == 0) r_norm = beta * r_norm;

      // compute new search direction
      if (tid == 0) pcps::xscal(n, beta, d, 1);
      if (tid == 0) pcps::xaxpy(n, pcps::unity<S>(), r, 1, d, 1);

    }

    // check if the method converged
    const bool converged = a->converged(std::abs(std::sqrt(r_norm)), myrank, tid);
    if (converged && myrank == 0 && tid == 0)
      std::cout << boost::format("conjugate gradient converged in %10i iterations") % iter << std::endl << std::endl;
    else if (myrank == 0 && tid == 0)
      std::cout << boost::format("conjugate gradient did not converge after %10i iterations") % iter << std::endl << std::endl;
      //throw pcps::Exception( (boost::format("pcps::conjugate_gradient did not converge.  residual = %.2e") % std::abs(std::sqrt(r_norm))).str() );

    // we currently have the solution to the equation (M A M) (M^(-1) x) = (M b), so multiply by M to get solution to A x = b
    a->operate_by_M(n, v, x, myrank, tid);

  }

} // end namespace pcps

#endif
